package indi.zhifa.expshare.class001.chemicalmatch.util.redis.impl;

import indi.zhifa.expshare.class001.chemicalmatch.properties.VectorRedisProperties;
import indi.zhifa.expshare.class001.chemicalmatch.util.openai.embding.OpenAiEmbedUtil;
import indi.zhifa.expshare.class001.chemicalmatch.util.redis.INickRedisUtil;
import indi.zhifa.engine.core.common.entity.web.ServiceException;
import jakarta.annotation.PostConstruct;
import lombok.RequiredArgsConstructor;
import lombok.extern.slf4j.Slf4j;
import org.apache.commons.codec.digest.MurmurHash3;
import org.apache.commons.lang3.ArrayUtils;
import org.apache.commons.pool2.impl.GenericObjectPoolConfig;
import org.springframework.stereotype.Component;
import redis.clients.jedis.*;
import redis.clients.jedis.search.Document;
import redis.clients.jedis.search.FTCreateParams;
import redis.clients.jedis.search.IndexDataType;
import redis.clients.jedis.search.Query;
import redis.clients.jedis.search.schemafields.SchemaField;
import redis.clients.jedis.search.schemafields.TagField;
import redis.clients.jedis.search.schemafields.TextField;
import redis.clients.jedis.search.schemafields.VectorField;

import java.nio.ByteBuffer;
import java.nio.ByteOrder;
import java.util.*;

@Slf4j
@RequiredArgsConstructor
@Component
public class NickNameRedisUtil implements INickRedisUtil {

    final OpenAiEmbedUtil mOpenAiEmbedUtil;
    final VectorRedisProperties mVectorRedisProperties;
    UnifiedJedis mJedis;
    static final String NICK_NAME_INDEX_NAME = "idx:nick_name";
    static final String NICK_NAME_KEY_PREFIX = "nick";
    static final int EMBEDDING_DIM = 1536;
    static final float DISTANCE_THRESHOLD = 0.6f;

    @PostConstruct
    public void init(){

        DefaultJedisClientConfig defaultJedisClientConfig = DefaultJedisClientConfig.builder()
                .clientName("jedis-vec")
                .database(mVectorRedisProperties.getDatabase())
                .password(mVectorRedisProperties.getPassword())
                .timeoutMillis((int)mVectorRedisProperties.getTimeout().toMillis())
                .connectionTimeoutMillis((int)mVectorRedisProperties.getTimeout().toMillis())
                .build();
        HostAndPort hostAndPort = new HostAndPort(mVectorRedisProperties.getHost(), mVectorRedisProperties.getPort());

        if(null != mVectorRedisProperties.getPool() && mVectorRedisProperties.getPool().getEnabled()){
            GenericObjectPoolConfig<Connection> poolConfig = new GenericObjectPoolConfig<Connection>();
            poolConfig.setMaxIdle(mVectorRedisProperties.getPool().getMaxIdle());
            poolConfig.setMinIdle(mVectorRedisProperties.getPool().getMinIdle());
            poolConfig.setMaxTotal(mVectorRedisProperties.getPool().getMaxActive());
            mJedis = new JedisPooled(poolConfig,hostAndPort,defaultJedisClientConfig);

        }else{
            mJedis = new UnifiedJedis(hostAndPort,defaultJedisClientConfig);
        }
    }


    @Override
    public void addIndex(String pModuleName, String pSubModuleName) {
        Set<String> indexSet = mJedis.ftList();
        String indexName = NICK_NAME_INDEX_NAME +":"+ pModuleName + ":" + pSubModuleName;
        String nickNamePrefix = NICK_NAME_KEY_PREFIX +":"+ pModuleName + ":" + pSubModuleName + ":";

        if(indexSet.contains(indexName)){
            return;
        }
        SchemaField[] schema = {
                TextField.of("nick_name"),
                TagField.of("name"),
                VectorField.builder()
                        .fieldName("embedding")
                        .algorithm(VectorField.VectorAlgorithm.HNSW)
                        .attributes(
                                Map.of(
                                        "TYPE", "FLOAT32",
                                        "DIM", EMBEDDING_DIM,
                                        "DISTANCE_METRIC", "L2"
                                )
                        )
                        .build()
        };
        mJedis.ftCreate(indexName,
                FTCreateParams.createParams()
                        .addPrefix(nickNamePrefix)
                        .on(IndexDataType.HASH),
                schema
        );
    }

    protected byte[] vec2ByteString(List<Float> vec) {
        byte[] bytes = new byte[vec.size() * Float.BYTES];

        Float[] boxed = vec.toArray(new Float[0]);
        float[] array = ArrayUtils.toPrimitive(boxed, 0f);

        ByteBuffer
                .wrap(bytes)
                .order(ByteOrder.LITTLE_ENDIAN)
                .asFloatBuffer()
                .put(array);

        return bytes;
    }

    @Override
    public void addNickName(String pModuleName, String pSubModuleName, String pName, List<String> pNickNames) {

        String nickNamePrefix = NICK_NAME_KEY_PREFIX +":"+ pModuleName + ":" + pSubModuleName + ":";
        List<List<Float>> embeddings = mOpenAiEmbedUtil.getEmbedding(pNickNames);

        for (int i=0;i<pNickNames.size();i++) {
            String nickName = pNickNames.get(i);
            int fingerPrint = MurmurHash3.hash32(nickName);
            String key = nickNamePrefix + fingerPrint;

            if(mJedis.exists(key)){
                return;
            }

            mJedis.hset(key,
                    Map.of(
                            "nick_name", nickName,
                            "name", pName
                    )
            );
            mJedis.hset(
                    key.getBytes(),
                    "embedding".getBytes(),
                    vec2ByteString(embeddings.get(i))
            );
            log.info("已添加缓存:"+key+"|  "+nickName+" -> " + pName);
        }

    }

    @Override
    public String queryNickName(String pModuleName, String pSubModuleName, String pNickName) {
        String indexName = NICK_NAME_INDEX_NAME +":"+ pModuleName + ":" + pSubModuleName;
        List<List<Float>> embeddings = mOpenAiEmbedUtil.getEmbedding(Arrays.asList(pNickName));
        if(embeddings.size() == 0){
            throw new ServiceException("获取"+pNickName+"向量失败");
        }
        byte[] bVec = vec2ByteString(embeddings.get(0));

        int K = 3;
        Query q = new Query("*=>[KNN $K @embedding $BLOB AS distance]")
                .returnFields("nick_name", "name", "distance")
                .addParam("K", K)
                .addParam(
                        "BLOB",
                        bVec
                )
                .setSortBy("distance", true)
                .dialect(2);

        List<Document> docs = mJedis.ftSearch(indexName, q).getDocuments();

        for (Document doc: docs) {
            System.out.println(
                    String.format(
                            "ID: %s, Distance: %s, nick_name: %s, name: %s",
                            doc.getId(),
                            doc.get("distance"),
                            doc.get("nick_name"),
                            doc.get("name")
                    )
            );
        }
        if(docs.size() == 0){
            return null;
        }
        String distance = docs.get(0).getString("distance");
        Float fDistance = Float.parseFloat(distance);
        if(fDistance > DISTANCE_THRESHOLD){
            return null;
        }

        return docs.get(0).getString("name");
    }

    @Override
    public void deleteNickName(String pModuleName, String pSubModuleName, String pName,
                               List<String> pNickNames) {
        String nickNamePrefix = NICK_NAME_KEY_PREFIX +":"+ pModuleName + ":" + pSubModuleName + ":";
        List<String> delKeys = new ArrayList<>();
        for (int i=0;i<pNickNames.size();i++) {
            String nickName = pNickNames.get(i);
            int fingerPrint = MurmurHash3.hash32(nickName);
            String key = nickNamePrefix + fingerPrint;
            delKeys.add(key);
        }
        mJedis.del(delKeys.toArray(String[]::new));
    }
}
